package com.baidu.acg.industry.bigdata.controller;


import com.baidu.acg.industry.bigdata.config.ApplicationConfiguration;
import com.baidu.acg.industry.bigdata.service.K8sDeployer;
import com.baidu.acg.industry.bigdata.service.YongJiModelService;
import io.swagger.annotations.Api;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

@Api(tags = "数据API")
@Slf4j
@RestController
@RequestMapping(value = "/v1", produces = "application/json;charset=UTF-8")
@RequiredArgsConstructor(onConstructor = @__(@Autowired))
public class ContainerController {

    private final ApplicationConfiguration modelConfiguration;
    private final YongJiModelService yongJiModelService;
    private final K8sDeployer k8sDeployer;

    @GetMapping("/health")
    public String getHealth() {
        return "ok";
    }

    @GetMapping("/create/pod")
    public String testModelRun() {
//        k8sDeployer.createK8sPod();
        return "success create pod";
    }

    @GetMapping("/model/train")
    public Map<String, String> getModelTrainFile() {
        Map<String, String> stringStringMap = yongJiModelService.modelTrain(modelConfiguration.getModelHostFile(), "document");
        log.info("模型下载路径=>" + stringStringMap.toString());
        return stringStringMap;
    }


    // Docker命令参数化模板
    public void executeDistributedTraining() throws Exception {
        List<String> command = new ArrayList<>();
        command.add("nvidia-docker");
        command.add("run");
        command.add("--rm");
        command.add("-d");
        command.add("--name");
        command.add("auto_train");
        command.add("-v");
        command.add("/ssd2:/workspace");
        command.add("--network=host");
        command.add("--shm-size");
        command.add("64g");
        command.add("-w");
        command.add("/workspace/hesensen/中车永济/PaddleScience/examples/_zhongche_gino");
        command.add("zc_paddle_test_env:cu118_py310");
        command.add("/bin/bash");
        command.add("-c");
        command.add("CUDA_VISIBLE_DEVICES=2,3 python3.10 -m paddle.distributed.launch --gpus='2,3' main.py");

        ProcessBuilder processBuilder = new ProcessBuilder(command);
        processBuilder.redirectErrorStream(true); // 合并标准错误与输出

        try {
            Process process = processBuilder.start();
            // 异步读取输出流（防止阻塞）
            new Thread(() -> {
                try (InputStream is = process.getInputStream();
                     BufferedReader reader = new BufferedReader(new InputStreamReader(is))) {
                    String line;
                    while ((line = reader.readLine()) != null) {
                        System.out.println("[Docker Log] " + line); // 可替换为日志框架
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }).start();

            // 设置超时与等待
            boolean completed = process.waitFor(2, TimeUnit.HOURS);
            if (!completed) {
                throw new RuntimeException("训练任务超时");
            }
            int exitCode = process.exitValue();
            if (exitCode != 0) {
                throw new RuntimeException("训练失败，退出码: " + exitCode);
            }
        } catch (Exception e) {
            throw new RuntimeException("Docker进程执行异常", e);
        }
    }


//    public static void main(String[] args) {
//        try {
//            new DockerController().executeDistributedTraining();
//        } catch (Exception e) {
//            e.printStackTrace();
//        }
//    }
}